package com.welson.db.cryptor.interceptor;

import com.google.common.collect.Lists;
import com.welson.db.cryptor.annotation.CryptField;
import com.welson.db.cryptor.config.DbCryptorProperties;
import com.welson.db.cryptor.core.EntityCryptProcessor;
import com.welson.db.cryptor.core.EntityCryptProcessorFactory;
import com.welson.db.cryptor.enums.CryptOperation;
import com.welson.db.cryptor.util.Cryptor;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.binding.MapperMethod;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.util.CollectionUtils;
import org.springframework.util.ReflectionUtils;

import java.lang.reflect.Method;
import java.lang.reflect.Parameter;
import java.util.*;

/**
 * 数据加解密拦截器
 * @author welson
 */
@Slf4j
@Intercepts({
        @Signature(type = Executor.class, method = "update", args = {MappedStatement.class, Object.class}),
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class})
})
public class CryptInterceptor implements Interceptor {
    private static final String GENERIC_NAME_PREFIX = "param";
    private final DbCryptorProperties properties;
    private final EntityCryptProcessorFactory processorFactory;

    public CryptInterceptor(DbCryptorProperties properties, EntityCryptProcessorFactory processorFactory) {
        this.properties = properties;
        this.processorFactory = processorFactory;
    }

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        if (!properties.isEnabled()) {
            return invocation.proceed();
        }

        Object[] args = invocation.getArgs();
        MappedStatement ms = (MappedStatement) args[0];
        Method mapperMethod = getMapperMethod(ms);

        // 处理入参加密
        handleParameterEncrypt(mapperMethod, args);

        // 执行原始方法
        Object result = invocation.proceed();

        // 处理结果解密
        handleResultDecrypt(mapperMethod, result);

        // 还原入参
        restoreParameter(args[1]);

        return result;
    }

    @SuppressWarnings("unchecked")
    private void restoreParameter(Object parameter) {
        if (parameter == null) {
            return;
        }
        // 处理单个实体类参数
        if (isEligibleEntityClass(parameter.getClass())) {
            processEntityDecryption(parameter);
        }
        // 处理update/updateById/deleteById方法的实体类参数
        if (parameter instanceof MapperMethod.ParamMap) {
            MapperMethod.ParamMap<Object> paramMap = (MapperMethod.ParamMap<Object>) parameter;
            Object entity = paramMap.getOrDefault("et", null);
            if (entity != null && isEligibleEntityClass(entity.getClass())) {
                processEntityDecryption(entity);
            }
        }
    }

    @SuppressWarnings("unchecked")
    private void handleParameterEncrypt(Method mapperMethod, Object[] args) {
        Object parameter = args[1];
        if (parameter == null || mapperMethod == null) {
            return;
        }

        Parameter[] parameters = mapperMethod.getParameters();

        if (parameter instanceof String) {
            if (shouldEncrypt(parameters[0])) {
                args[1] = Cryptor.encrypt((String) parameter);
            }
        } else if (parameter instanceof MapperMethod.ParamMap) {
            handleParamMapEncrypt((MapperMethod.ParamMap<Object>) parameter, parameters);
        } else if (parameter instanceof Map) {
            handleMapEncrypt((Map<Object, Object>) parameter);
        } else {
            processEntityEncryption(parameter);
        }
    }

    private void handleParamMapEncrypt(MapperMethod.ParamMap<Object> paramMap, Parameter[] parameters) {
        Set<Object> processedSet = new HashSet<>();

        // 处理字符串参数
        for (int i = 0; i < parameters.length; i++) {
            Parameter param = parameters[i];
            if (shouldEncrypt(param)) {
                String paramName = getParamName(param);
                Object paramValue = paramMap.get(paramName);
                if (paramValue == null) {
                    continue;
                }
                processParamValue(paramMap, paramValue, paramName, i, parameters.length, processedSet);
            }
        }

        // 处理实体类参数
        paramMap.forEach((key, value) -> {
            if (value != null && !processedSet.contains(value)) {
                processEntityEncryption(value);
                processedSet.add(value);
            }
        });
    }

    private void processParamValue(MapperMethod.ParamMap<Object> paramMap, Object paramValue, String paramName, int index, int paramCount, Set<Object> processedSet) {
        if (paramValue instanceof String) {
            Object encryptedValue = Cryptor.encrypt((String) paramValue);
            paramMap.put(paramName, encryptedValue);
            String genericName = GENERIC_NAME_PREFIX + (index + 1);
            paramMap.put(genericName, encryptedValue);
            processedSet.add(paramValue);
        } else if (paramValue instanceof Collection) {
            Collection<?> collection = (Collection<?>) paramValue;
            if (!CollectionUtils.isEmpty(collection) && collection.iterator().next() instanceof String) {
                Collection<Object> encryptedCollection = encryptCollection(collection);
                paramMap.put(paramName, encryptedCollection);
                if (paramCount == 1) {
                    paramMap.put("collection", encryptedCollection);
                    if (paramValue instanceof List) {
                        paramMap.put("list", encryptedCollection);
                    }
                } else {
                    String genericName = GENERIC_NAME_PREFIX + (index + 1);
                    paramMap.put(genericName, encryptedCollection);
                }
                processedSet.add(paramValue);
            }
        } else if (paramValue.getClass().isArray()) {
            Object[] array = (Object[]) paramValue;
            Object[] encryptedArray = new Object[array.length];
            for (int i = 0; i < array.length; i++) {
                if (array[i] instanceof String) {
                    encryptedArray[i] = Cryptor.encrypt((String) array[i]);
                } else {
                    encryptedArray[i] = array[i];
                }
            }
            paramMap.put(paramName, encryptedArray);
            if (paramCount == 1) {
                paramMap.put("array", encryptedArray);
            } else {
                String genericName = GENERIC_NAME_PREFIX + (index + 1);
                paramMap.put(genericName, encryptedArray);
            }
            processedSet.add(paramValue);
        }
    }

    private void handleMapEncrypt(Map<Object, Object> map) {
        map.forEach((key, value) -> {
            if (value instanceof Collection) {
                Collection<?> collection = (Collection<?>) value;
                if (!CollectionUtils.isEmpty(collection)) {
                    processEntityEncryption(collection);
                }
            } else if (value.getClass().isArray()) {
                handleArrayEncrypt(value);
            } else {
                processEntityEncryption(value);
            }
        });
    }

    private void handleArrayEncrypt(Object array) {
        if (array == null) {
            return;
        }
        Object[] objects = (Object[]) array;
        for (int i = 0; i < objects.length; i++) {
            if (objects[i] instanceof String) {
                objects[i] = Cryptor.encrypt((String) objects[i]);
            } else {
                processEntityEncryption(objects[i]);
            }
        }
    }

    private void handleResultDecrypt(Method mapperMethod, Object result) {
        if (result == null || mapperMethod == null) {
            return;
        }

        CryptField methodCryptField = mapperMethod.getAnnotation(CryptField.class);
        if (shouldDecrypt(methodCryptField)) {
            if (result instanceof String) {
                return;
            } else if (result instanceof Collection) {
                Collection<?> collection = (Collection<?>) result;
                if (!CollectionUtils.isEmpty(collection) && collection.iterator().next() instanceof String) {
                    decryptCollection(collection);
                    return;
                }
            } else if (result.getClass().isArray()) {
                handleArrayDecrypt(result);
                return;
            }
        }

        processEntityDecryption(result);
    }

    private void handleArrayDecrypt(Object array) {
        if (array == null) {
            return;
        }
        Object[] objects = (Object[]) array;
        for (int i = 0; i < objects.length; i++) {
            if (objects[i] instanceof String) {
                objects[i] = Cryptor.decrypt((String) objects[i]);
            } else {
                processEntityDecryption(objects[i]);
            }
        }
    }

    private Method getMapperMethod(MappedStatement ms) {
        try {
            String id = ms.getId();
            String className = id.substring(0, id.lastIndexOf("."));
            String methodName = id.substring(id.lastIndexOf(".") + 1);
            Class<?> mapperClass = Class.forName(className);

            Method[] allDeclaredMethods = ReflectionUtils.getAllDeclaredMethods(mapperClass);
            for (Method method : allDeclaredMethods) {
                if (method.getName().equals(methodName)) {
                    return method;
                }
            }
        } catch (Exception e) {
            log.error("Failed to get mapper method", e);
        }
        return null;
    }

    private String getParamName(Parameter parameter) {
        org.apache.ibatis.annotations.Param param = parameter.getAnnotation(org.apache.ibatis.annotations.Param.class);
        return param != null ? param.value() : parameter.getName();
    }

    private void processEntityEncryption(Object entity) {
        if (entity == null) {
            return;
        }

        if (entity instanceof Collection) {
            ((Collection<?>) entity).forEach(this::processEntityEncryption);
            return;
        }

        if (entity instanceof Map) {
            Map<?, ?> map = (Map<?, ?>) entity;
            for (Object value : map.values()) {
                processEntityEncryption(value);
            }
            return;
        }

        if (entity.getClass().isArray()) {
            handleArrayEncrypt(entity);
            return;
        }

        EntityCryptProcessor<Object> processor = processorFactory.getProcessor(entity);
        processor.processEncrypt(entity);
    }

    private void processEntityDecryption(Object entity) {
        if (entity == null) {
            return;
        }

        if (entity instanceof Collection) {
            ((Collection<?>) entity).forEach(this::processEntityDecryption);
            return;
        }

        if (entity instanceof Map) {
            Map<?, ?> map = (Map<?, ?>) entity;
            for (Object value : map.values()) {
                processEntityDecryption(value);
            }
            return;
        }

        if (entity.getClass().isArray()) {
            handleArrayDecrypt(entity);
            return;
        }


        EntityCryptProcessor<Object> processor = processorFactory.getProcessor(entity);
        processor.processDecrypt(entity);
    }

    private Collection<Object> encryptCollection(Collection<?> collection) {
        if (CollectionUtils.isEmpty(collection)) {
            return Collections.emptyList();
        }

        Collection<Object> encryptedCollection;
        if (collection instanceof List) {
            encryptedCollection = new ArrayList<>(collection.size());
        } else if (collection instanceof Set) {
            encryptedCollection = new HashSet<>(collection.size());
        } else {
            return Lists.newArrayList(collection);
        }

        for (Object item : collection) {
            if (item instanceof String) {
                encryptedCollection.add(Cryptor.encrypt((String) item));
            } else {
                encryptedCollection.add(item);
            }
        }

        return encryptedCollection;
    }

    @SuppressWarnings("unchecked")
    private void decryptCollection(Collection<?> collection) {
        if (CollectionUtils.isEmpty(collection)) {
            return;
        }

        Collection<Object> decryptedCollection;
        if (collection instanceof List) {
            decryptedCollection = new ArrayList<>(collection.size());
        } else if (collection instanceof Set) {
            decryptedCollection = new HashSet<>(collection.size());
        } else {
            return;
        }

        for (Object item : collection) {
            if (item instanceof String) {
                decryptedCollection.add(Cryptor.decrypt((String) item));
            } else {
                decryptedCollection.add(item);
            }
        }

        collection.clear();
        ((Collection<Object>) collection).addAll(decryptedCollection);
    }

    private boolean shouldDecrypt(CryptField cryptField) {
        if (cryptField == null) {
            return false;
        }
        return cryptField.operation() == CryptOperation.DECRYPT || cryptField.operation() == CryptOperation.BOTH;
    }

    private boolean shouldEncrypt(Parameter parameter) {
        if (parameter == null) {
            return false;
        }
        CryptField annotation = parameter.getAnnotation(CryptField.class);
        return annotation != null && (annotation.operation() == CryptOperation.ENCRYPT || annotation.operation() == CryptOperation.BOTH);
    }

    private boolean isEligibleEntityClass(Class<?> entityClass) {
        if (entityClass == null) {
            return false;
        }
        String className = entityClass.getName();
        return properties.getScanPackages().stream().anyMatch(className::startsWith);
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {
        // 可以从配置中读取加解密相关的属性
    }
} 